from sys import argv
from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GLUT import *

class MyPyOpenGLTest:
    def __init__(self, width=640, height=480, title='画刷'):
        glutInit(len(argv), argv)
        glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
        glutInitWindowSize(width, height)
        glutInitWindowPosition(350, 200)
        self.window = glutCreateWindow(title.encode('gbk'))
        glutDisplayFunc(self.draw)
        glutReshapeFunc(self.reshape)
        glutMouseFunc(self.mouse_downup)
        glutMotionFunc(self.mouse_move)
        glutKeyboardFunc(self.ascii_key)
        glClearColor(1, 1, 1, 1)

        # 创建菜单
        self.create_menu()
        # 初始化颜色和线宽
        self.data = [(0, 0, 0), (1,)]
        # 标记鼠标左键是否按下
        self.drawing = False
        # 绘制模式， 不是橡皮擦模式
        self.is_eraser = False

    def process_menu(self, value):
        div, mod = divmod(value, 10)
        if div == 1:
            self.data.append((mod,))
        elif div == 2:
            colors = {1: (1, 0, 0), 2: (0, 1, 0), 3: (0, 0, 1)}
            if mod in (1, 2, 3):
                self.data.append(colors[mod])
                self.is_eraser = False
        elif div == 3:
            # 清除图形
            del self.data[2:]
            glutPostRedisplay()
        return value

    def create_menu(self):
        # 创建子菜单
        sub_menu1 = glutCreateMenu(self.process_menu)
        # 数值第一位表示菜单功能， 第二位与线宽对应
        glutAddMenuEntry('1 pixel', 11)
        glutAddMenuEntry('3 pixel', 13)
        glutAddMenuEntry('5 pixel', 15)

        sub_menu2 = glutCreateMenu(self.process_menu)
        # 数值第二位与颜色对应
        glutAddMenuEntry('red', 21)
        glutAddMenuEntry('green', 22)
        glutAddMenuEntry('blue', 23)
        # 创建主菜单，添加子菜单
        main_menu = glutCreateMenu(self.process_menu)
        glutAddSubMenu('Line Width', sub_menu1)
        glutAddSubMenu('Line Color', sub_menu2)
        glutAddMenuEntry('Reset', 30)
        glutAttachMenu(GLUT_RIGHT_BUTTON)


    def mouse_downup(self, button, state, x, y):
        if button != GLUT_LEFT_BUTTON:
            # 只响应鼠标左键，其他按键直接返回
            return
        if state == GLUT_DOWN:
            self.drawing = True
        elif state == GLUT_UP:
            self.drawing = False
            # None表示一段折线绘制结束
            self.data.append(None)

    def mouse_move(self, x, y):
        if self.drawing:
            # 追加当前位置，重新绘制图形
            self.data.append((x, y))
            glutPostRedisplay()

    def ascii_key(self, key, x, y):
        if key == b' ':
            last_two_colors = tuple(filter(lambda item: item and len(item) - -3, self.data))[-2:]
            if last_two_colors[-1] == (1, 1, 1):
                # 重新设置为最后一次设置的背景颜色
                self.data.append(last_two_colors[-2])
            else:
                # 背景色， 白色
                self.data.append((1, 1, 1))
            self.is_eraser = not self.is_eraser

    def reshape(self, w, h):
        glViewport(0, 0, w, h)
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        gluOrtho2D(0, w, h, 0)
        glMatrixMode(GL_MODELVIEW)
        glLoadIdentity()
        
    def draw(self):
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        glMatrixMode(GL_MODELVIEW)
        glLoadIdentity()
        # 绘制折线图
        glBegin(GL_LINE_STRIP)
        for item in self.data:
            if item is None:
                glEnd()
                glBegin(GL_LINE_STRIP)
            else:
                item_len = len(item)
                if item_len == 1:
                    # 结束上一段折线，修改线宽，开始绘制下一段折线
                    glEnd()
                    glLineWidth(item[0])
                    glBegin(GL_LINE_STRIP)
                elif item_len == 2:
                    glVertex2fv(item)
                elif item_len == 3:
                    # 结束上一段折线，修改颜色，开始绘制下一段折线
                    glEnd()
                    glColor3fv(item)
                    glBegin(GL_LINE_STRIP)
        glEnd()
        
        glutSwapBuffers()

    # 消息主循环
    def mainloop(self):
        glutMainLoop()

if __name__ == '__main__':
    MyPyOpenGLTest().mainloop()
